import {
  any
} from "../../../core/logic";
import {
  Compute
} from "./Compute";
import {
  gtao,
  gtao_multisampled
} from "./shaders/gtao";
import {
  Pass
} from "../Pass";
import {
  Engine
} from "../../../Engine";
import {
  guid
} from "../../../core/utils";
/**
 * @class
 * @memberof w.compute
 */
class PickComputer extends Compute {
  constructor(opts = {}) {
    super(opts);
    this.initBuffer()
  }
  initBuffer() {
    this.xyBuffer = Engine.instance.device.createBuffer({
      label: "xy buffer",
      size: 16,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
    });
    this.outBuffer = Engine.instance.device.createBuffer({
      label: "outBuffer",
      size: 4 * 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    });
    this.outBuffer1 = Engine.instance.device.createBuffer({
      label: "out1Buffer",
      size: 4 * 4,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
    });
    this.readBuffer = Engine.instance.device.createBuffer({
      size: this.outBuffer.size,
      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    });
    this.readBuffer1 = Engine.instance.device.createBuffer({
      size: this.outBuffer.size,
      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
    });
  }
  getShader() {
    const sampleLoop = () => {
      if (Engine.instance.enableGbuffer) {
        return `
               color = textureLoad(colorTexture, vec2<u32>(u32(xy.x), u32(xy.y)), 0);  
                position += textureLoad(positionTexture, vec2<u32>(u32(xy.x), u32(xy.y)),0); 
               `
      }
      if (!Engine.instance.enableGbuffer) {
        return `
               color = textureLoad(colorTexture, vec2<u32>(u32(xy.x), u32(xy.y)), 0);  
               `
      }
    }
    return `  
        @group(0) @binding(0) var colorTexture: texture_2d<f32>;  
       ${Engine.instance.enableGbuffer ? '@group(0) @binding(1) var positionTexture: texture_2d<f32>;' : ''}   
        @group(0) @binding(2) var<uniform> xy: vec2<f32>;  
        @group(0) @binding(3) var<storage, read_write> outBuffer: vec4<f32>;  
        ${Engine.instance.enableGbuffer ? '@group(0) @binding(4) var<storage, read_write> outBuffer1: vec4<f32>;' : ''}  

        @compute @workgroup_size(1)  
        fn main() {  
            var color: vec4<f32> = vec4<f32>(0.0);  
            var position: vec4<f32> = vec4<f32>(0.0); 
            ${sampleLoop()}
            outBuffer = color;  
            ${Engine.instance.enableGbuffer ? 'outBuffer1 = position;' : ''}  
        }  
    `;
  }

  initPipeLine() {
    const id = Engine.instance.enableGbuffer + Engine.instance.multisample
    if (!this.computePipeline && (this.computePipeline?.id != id)) {
      const shaderModule = Engine.instance.device.createShaderModule({
        code: this.getShader(),
      });
      // 创建计算管线
      this.computePipeline = Engine.instance.device.createComputePipeline({
        label: "pick computer",
        layout: `auto`,
        compute: {
          module: shaderModule,
          entryPoint: 'main',
        },
      });
      this.computePipeline.id = Engine.instance.enableGbuffer + Engine.instance.multisample
    }
  }
  getBindGroup(enableGbuffer, colorTexture, positionTexture) {
    if (enableGbuffer) {
      return Engine.instance.device.createBindGroup({
        label: "pick group",
        layout: this.computePipeline.getBindGroupLayout(0),
        entries: [
          { binding: 0, resource: colorTexture.view },
          { binding: 1, resource: positionTexture.view },
          {
            binding: 2, resource: {
              buffer: this.xyBuffer,
            },
          },
          {
            binding: 3, resource: {
              buffer: this.outBuffer,
            },
          },
          {
            binding: 4, resource: {
              buffer: this.outBuffer1,
            },
          },
        ],
      });
    } else {
      return Engine.instance.device.createBindGroup({
        label: "pick group",
        layout: this.computePipeline.getBindGroupLayout(0),
        entries: [
          { binding: 0, resource: colorTexture.view },
          {
            binding: 2, resource: {
              buffer: this.xyBuffer,
            },
          },
          {
            binding: 3, resource: {
              buffer: this.outBuffer,
            },
          },
        ],
      });
    }
  }
  initBindGroup(colorTexture, positionTexture) {
    // 创建绑定组
    if (!this.bindGroup || (this.bindGroup.id != colorTexture.id)) {
      this.bindGroup = this.getBindGroup(Engine.instance.enableGbuffer, colorTexture, positionTexture)
      this.bindGroup.id = colorTexture.id
    }
  }
  async get(xydata, colorTexture, positionTexture) {
    if (!Engine.instance.enablePick) {
      return
    }
    if (Engine.instance.enablePick && !colorTexture) {
      return
    }
    if (Engine.instance.enablePick && colorTexture && (Engine.instance.enableGbuffer && !positionTexture)) {
      return
    }
    this.commandEncoder = Engine.instance.device.createCommandEncoder()
    this.initPipeLine(colorTexture)
    this.initBindGroup(colorTexture, positionTexture)
    Engine.instance.queue.writeBuffer(this.xyBuffer, 0, xydata);
    const computePass = this.commandEncoder.beginComputePass();
    computePass.setPipeline(this.computePipeline);
    computePass.setBindGroup(0, this.bindGroup);
    computePass.dispatchWorkgroups(1);
    computePass.end();
    this.commandEncoder.copyBufferToBuffer(this.outBuffer, 0, this.readBuffer, 0, this.outBuffer.size);
    this.commandEncoder.copyBufferToBuffer(this.outBuffer1, 0, this.readBuffer1, 0, this.outBuffer1.size);
    const commands = this.commandEncoder.finish({ label: "pick" });
    Engine.instance.queue.submit([commands]);
    await this.readBuffer.mapAsync(GPUMapMode.READ);
    await this.readBuffer1.mapAsync(GPUMapMode.READ);
    const mappedArray = this.readBuffer.getMappedRange();
    const mappedArray1 = this.readBuffer1.getMappedRange();
    let color = new Float32Array(mappedArray);
    let position = new Float32Array(mappedArray1);
    position = [position[0], position[1], position[2]]
    color = Array.from(color)
    position = Array.from(position)
    this.readBuffer?.unmap();
    this.readBuffer1?.unmap();
    return { color, position }
  }
}
export {
  PickComputer
};